package ai.chatgpt;

import cn.hutool.http.ContentType;
import cn.hutool.http.Header;
import cn.hutool.http.HttpRequest;
import com.alibaba.fastjson.JSON;
import com.fasterxml.jackson.annotation.JsonProperty;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;

import java.util.List;

/**
 * @author zhanglu
 * @project JavaBase
 * @description
 * @date 2023/2/8 23:35
 */
@Slf4j
@NoArgsConstructor
@Data
public class OpenAiApiTest {
    @JsonProperty("id")
    private String id;
    @JsonProperty("object")
    private String object;
    @JsonProperty("created")
    private Integer created;
    @JsonProperty("model")
    private String model;
    @JsonProperty("choices")
    private List<ChoicesDTO> choices;
    @JsonProperty("usage")
    private UsageDTO usage;

    @NoArgsConstructor
    @Data
    public static class UsageDTO {
        @JsonProperty("prompt_tokens")
        private Integer promptTokens;
        @JsonProperty("completion_tokens")
        private Integer completionTokens;
        @JsonProperty("total_tokens")
        private Integer totalTokens;
    }

    @NoArgsConstructor
    @Data
    public static class ChoicesDTO {
        @JsonProperty("text")
        private String text;
        @JsonProperty("index")
        private Integer index;
        @JsonProperty("logprobs")
        private Object logprobs;
        @JsonProperty("finish_reason")
        private String finishReason;
    }


    public static void main(String[] args) throws Exception {
        String params = "{\"model\": \"text-davinci-003\", \"prompt\": \"帮我写一个shell脚本，递归查询文件\", \"temperature\": 0, \"max_tokens\": 4000}";
        log.info("请求参数：{}", params);
        String res = HttpRequest.post("https://api.openai.com/v1/completions")
                .header(Header.CONTENT_TYPE, ContentType.JSON.getValue())
                .header("Authorization", "Bearer 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")
                .body(params)
                .timeout(8 * 60 * 1000)
                .execute().body();
        log.info("请求响应：{}", res);
        OpenAiApiTest openAiApiTest = JSON.parseObject(res, OpenAiApiTest.class);
        openAiApiTest.getChoices().forEach(o -> System.out.println(o.getText()));
    }

}
